# 更新模型的基础信息
import os
import sys
import glob
import shutil
from mmengine import Config
from mmdet.registry import MODELS
from mmdet.utils import register_all_modules
import argparse

register_all_modules()

models_dir = os.path.abspath(os.path.join(__file__, *(['..'] * 3)))
root_dir = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
sys.path.insert(0, root_dir)
from _base_.dataset_utils import DatasetConfigLoader
from _base_.model_arch import ModelParser
from _base_.model_info import ModelInspector

# 默认的数据集配置文件
balloon_dataset_cfg = """
dataset_type: COCO
classes: [balloon]

TrainDataset:
    dataset_dir: %s
    image_dir: val
    anno_path: val/annotation_coco.json

EvalDataset:
    dataset_dir: %s
    image_dir: val
    anno_path: val/annotation_coco.json

TestDataset:
    dataset_dir: %s
    image_dir: val
    anno_path: val/annotation_coco.json
"""

# 输入模型的目录，例如faster_rcnn，对应在config目标中
# 输出模型每一个配置文件的信息
def updateModelInfo( model_name ):
    # 创建保存配置文件的目录
    cfgs_dir = os.path.join( root_dir, model_name, 'model_infos' )
    if os.path.exists( cfgs_dir ):
        shutil.rmtree( cfgs_dir )
    os.makedirs(cfgs_dir)
    os.makedirs(os.path.join(cfgs_dir, 'arch'))
    os.makedirs(os.path.join(cfgs_dir, 'params'))

    # 遍历对应模型的所有配置文件
    model_dir = os.path.join(root_dir, 'configs', model_name)
    py_files = glob.glob(os.path.join(model_dir, '*.py'))
    model_parser = ModelParser()

    # 读取数据集配置
    dataset_dir = os.path.join( models_dir, 'data', 'balloon' )
    dataset_cfg = balloon_dataset_cfg % ( dataset_dir, dataset_dir, dataset_dir )
    config_loader = DatasetConfigLoader()
    inspector = ModelInspector()
    num_classes, dataset_cfg_dict = config_loader.load_yml_data(dataset_cfg)

    # 对应每一个配置文件，生成模型的结构和参数量
    cfg_list = []
    for cfg_path in py_files:
        print('## ' + cfg_path)
        try:
            cfg = Config.fromfile( cfg_path )
            model = MODELS.build( cfg.model )

            # 保存模型结构到对应文件（转换mmdetection模型结构）
            model_info = model_parser.parse(str(model))
            basename = os.path.splitext(os.path.basename(cfg_path))[0]
            model_arch_path = os.path.join( cfgs_dir, 'arch', basename + '.json')
            cfg_list.append( basename )
            print('处理模型[ %s ]' % basename)
            with open( model_arch_path, 'wt' ) as fd:
                fd.write( model_parser.to_json(model_info))

            # 保存模型参数量到对应文件（从torchinfo输出的结果）
            json_summary = inspector.inspect(cfg_path, dataset_cfg_dict)
            model_params_path = os.path.join( cfgs_dir, 'params', basename + '.json')
            with open( model_params_path, 'wt' ) as fd:
                fd.write( json_summary )
        except:
            print('处理模型[ %s ]失败' % cfg_path)

    cfg_list_path = os.path.join( cfgs_dir, 'model.list' )
    with open(cfg_list_path, 'wt') as fd:
        fd.write( '\n'.join(cfg_list) )

def getModelArch(cfg_path):
    # 保存模型结构到对应文件（转换mmdetection模型结构）
    model_parser = ModelParser()
    cfg = Config.fromfile( cfg_path )
    model = MODELS.build( cfg.model )

    # 保存模型结构到对应文件（转换mmdetection模型结构）
    model_info = model_parser.parse(str(model))  
    model_arch = model_parser.to_json(model_info)
    print(model_arch)

def main():
    parser = argparse.ArgumentParser(description='Script to update model info or run other functions.')
    parser.add_argument('-d', '--directory', type=str, help='Execute updateModelInfo with a given model name')
    parser.add_argument('-f', '--file', type=str, help='Execute another function with a given config file path')

    args = parser.parse_args()

    if args.directory:
        updateModelInfo(args.directory)
    elif args.file:
        getModelArch(args.file)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()